{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Augment LLM Basic Components\n",
        "\n",
        "These are just basic implementation examples for an Augmented LLM. Research and choose the best suited when building AI Agents and Workflows"
      ],
      "metadata": {
        "id": "2HJakOBhhE8I"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![image.png]()"
      ],
      "metadata": {
        "id": "n1ldbfzbhBqu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "6fOOS2iWTp3H"
      },
      "outputs": [],
      "source": [
        "%%capture --no-stderr\n",
        "%pip install --quiet -U langgraph langchain_google_genai python-dotenv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "WVbTUQ9q0m1d"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "os.environ[\"GOOGLE_API_KEY\"] = userdata.get('GOOGLE_API_KEY')\n",
        "\n",
        "os.environ[\"LANGCHAIN_API_KEY\"] = userdata.get('LANGCHAIN_API_KEY')\n",
        "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
        "os.environ[\"LANGCHAIN_PROJECT\"] = \"langchain-functional-api\""
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_google_genai import ChatGoogleGenerativeAI\n",
        "\n",
        "model = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash-exp\")"
      ],
      "metadata": {
        "id": "3EtakjdImZQd"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1. Structured Output"
      ],
      "metadata": {
        "id": "Ny7wLysb8ef2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pydantic import BaseModel, Field\n",
        "from langchain_google_genai import ChatGoogleGenerativeAI\n",
        "\n",
        "# Define output schema\n",
        "class SearchQuery(BaseModel):\n",
        "    search_query: str = Field(None, description=\"Query optimized for web search\")\n",
        "    justification: str = Field(None, description=\"Why this query is relevant\")\n",
        "\n",
        "# Create base LLM\n",
        "llm = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash-exp\")\n",
        "\n",
        "# Augment with structured output\n",
        "structured_llm = llm.with_structured_output(SearchQuery)\n",
        "\n",
        "# Use the augmented LLM\n",
        "result = structured_llm.invoke(\"How does Calcium CT score relate to high cholesterol?\")\n",
        "result"
      ],
      "metadata": {
        "id": "NrLQ3MUgH3z9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "072ad7d7-bdb4-4462-aadc-ee056206eb14"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "SearchQuery(search_query='calcium CT score high cholesterol relationship', justification=None)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2. Tool Usage"
      ],
      "metadata": {
        "id": "t1BzEQLSWDfG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def search_medical_database(query: str) -> str:\n",
        "    \"\"\"Search medical literature database\"\"\"\n",
        "    # Mock implementation\n",
        "    return f\"Results for: {query}\"\n",
        "\n",
        "# Bind tool to LLM\n",
        "llm_with_tools = llm.bind_tools([search_medical_database])\n",
        "\n",
        "# Use tool-augmented LLM\n",
        "response = llm_with_tools.invoke(\"Find recent studies on statins\")\n",
        "response.tool_calls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yTSnHdnZWDNH",
        "outputId": "4e29e9f4-c332-472c-c567-52786d4471a0"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'name': 'search_medical_database',\n",
              "  'args': {'query': 'recent studies on statins'},\n",
              "  'id': 'f88c1cbd-750e-4475-bf71-52228d8acf2a',\n",
              "  'type': 'tool_call'}]"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. Memory Integration\n"
      ],
      "metadata": {
        "id": "nP9MSbGDWG8M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import uuid\n",
        "from langgraph.store.memory import InMemoryStore, BaseStore\n",
        "from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings\n",
        "from langchain_core.runnables import RunnableConfig\n",
        "from langchain_core.messages import BaseMessage\n",
        "from langgraph.func import entrypoint, task\n",
        "from langgraph.graph import add_messages\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "\n",
        "\n",
        "# ✅ **Initialize LLM**\n",
        "model = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash-exp\")\n",
        "\n",
        "# ✅ **Initialize Memory Store for Bob-Agent**\n",
        "in_memory_store = InMemoryStore(\n",
        "    index={\n",
        "        \"embed\": GoogleGenerativeAIEmbeddings(model=\"models/embedding-001\"),  # Ensure correct model\n",
        "        \"dims\": 768,  # Ensure embedding dimensions match\n",
        "    }\n",
        ")\n",
        "\n",
        "\n",
        "# ✅ **Task: Bob-Agent Retrieves Its Own Memory**\n",
        "@task\n",
        "def call_model(messages: list[BaseMessage], memory_store: BaseStore):\n",
        "    \"\"\"Retrieves Bob-Agent's context from memory and responds accordingly.\"\"\"\n",
        "\n",
        "    namespace = (\"memories\", \"bob_agent\")\n",
        "\n",
        "    # 🔥 **Retrieve stored memories**\n",
        "    stored_memories = memory_store.search(namespace, query=\"Bob AI Agent persona and improvements\")\n",
        "    stored_info = \"\\n\".join([d.value[\"data\"] for d in stored_memories]) if stored_memories else \"\"\n",
        "    print(\"stored_info\", stored_info, \"\\n\\n\")\n",
        "\n",
        "    # ✅ **Adjust system prompt dynamically**\n",
        "    if stored_info:\n",
        "        system_msg = f\"You are Bob, a helpful AI assistant. Your current persona and improvements:\\n{stored_info}\"\n",
        "    else:\n",
        "        system_msg = \"You are Bob, a helpful AI assistant. No specific improvements or persona data is available yet.\"\n",
        "\n",
        "    # ✅ **Invoke the model with memory-augmented context**\n",
        "    response = model.invoke([{\"role\": \"system\", \"content\": system_msg}] + messages)\n",
        "\n",
        "    return response\n",
        "\n",
        "\n",
        "# ✅ **Task: Bob-Agent Learns from User Interactions**\n",
        "@task\n",
        "def learn_from_user(messages: list[BaseMessage], memory_store: BaseStore):\n",
        "    \"\"\"Stores learnings from a user interaction in memory.\"\"\"\n",
        "\n",
        "    namespace = (\"memories\", \"bob_agent\")\n",
        "\n",
        "    # 🔥 **Extract key learnings using LLM**\n",
        "    learnings_response = model.invoke([\n",
        "        {\"role\": \"system\", \"content\": \"Analyze the user chat and extract key improvements for yourself. Focus on persona refinement, response tone, and knowledge gaps. Do not store anything related to user.\"}\n",
        "    ] + messages)\n",
        "\n",
        "    # ✅ **Store learnings in Bob-Agent’s memory**\n",
        "    new_learning = f\"Bob-Agent learned: {learnings_response.content}\"\n",
        "    memory_store.put(namespace, str(uuid.uuid4()), {\"data\": new_learning})\n",
        "\n",
        "    return new_learning\n",
        "\n",
        "\n",
        "# ✅ **Workflow EntryPoint: Handles Bob-Agent's Persistent Memory**\n",
        "@entrypoint(checkpointer=MemorySaver(), store=in_memory_store)\n",
        "def workflow(\n",
        "    inputs: list[BaseMessage],\n",
        "    *,\n",
        "    previous: list[BaseMessage],\n",
        "    config: RunnableConfig,\n",
        "    store: BaseStore,\n",
        "):\n",
        "    \"\"\"Handles Bob-Agent's conversation flow and learning.\"\"\"\n",
        "\n",
        "    previous = previous or []\n",
        "    inputs = add_messages(previous, inputs)\n",
        "\n",
        "    # 🔥 **Retrieve response based on Bob-Agent's memory**\n",
        "    response = call_model(inputs, store).result()\n",
        "\n",
        "    # 🔥 **Bob-Agent learns from the interaction**\n",
        "    learn_from_user(inputs, store).result()\n",
        "\n",
        "    return entrypoint.final(value=response, save=add_messages(inputs, response))\n"
      ],
      "metadata": {
        "id": "mvKXn5YyWHal"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# ✅ **Example Usage: Testing Bob-Agent’s Learning & Memory**\n",
        "config1 = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"user_123\"}}\n",
        "input_message1 = {\"role\": \"user\", \"content\": \"LangGraph and CrewAI are great Agentic Framrworkd.\"}\n",
        "\n",
        "for chunk in workflow.stream([input_message1], config1, stream_mode=\"values\"):\n",
        "    chunk.pretty_print()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YydCaGwnWP7s",
        "outputId": "2dad0add-56d8-4589-b35c-8ea147580d46"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "stored_info  \n",
            "\n",
            "\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "You're right! LangGraph and CrewAI are both excellent frameworks for building agentic applications, and they offer different strengths. It's great to see the ecosystem evolving with different approaches to this exciting area.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# ✅ **Retrieve Bob-Agent's learned memory in a new session**\n",
        "config2 = {\"configurable\": {\"thread_id\": \"3\", \"user_id\": \"user_1234\"}}\n",
        "input_message2 = {\"role\": \"user\", \"content\": \"Can you share any good Agentic Framrwork\"}\n",
        "\n",
        "for chunk in workflow.stream([input_message2], config2, stream_mode=\"values\"):\n",
        "    chunk.pretty_print()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bDdtd1c2Xf2W",
        "outputId": "f33c5e54-ffbd-49a0-f97b-a8a4bf3eb80d"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "stored_info Bob-Agent learned: Okay, I understand. Based on that short statement, here's what I can consider for improvement:\n",
            "\n",
            "*   **Persona Refinement:** The user seems knowledgeable about agentic frameworks. I should avoid overly basic explanations when discussing LangGraph and CrewAI. I can assume a certain level of familiarity with concepts like agents, tools, and orchestration.\n",
            "*   **Response Tone:** A simple acknowledgement is a good start, but I can aim for a more engaging response. I could show enthusiasm for the topic or ask a follow-up question to encourage further conversation.\n",
            "*   **Knowledge Gaps:** This tells me I *should* have a solid understanding of both LangGraph and CrewAI. If asked specific questions about their differences, strengths, weaknesses, or use cases, I need to be prepared to answer accurately and comprehensively. I should ensure I'm up-to-date on the latest developments in these frameworks. \n",
            "\n",
            "\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "Of course! When it comes to agentic frameworks, two popular and powerful choices are LangGraph and CrewAI.\n",
            "\n",
            "To help me recommend the best one for you, could you tell me a bit more about what you're hoping to achieve with it? For example:\n",
            "\n",
            "*   What kind of tasks will the agents be performing?\n",
            "*   Do you need a highly customizable graph-based approach, or something more streamlined for collaborative tasks?\n",
            "*   What's your level of experience with agentic frameworks?\n",
            "\n",
            "In the meantime, I can give you a brief overview of each:\n",
            "\n",
            "*   **LangGraph:** This is a framework built on top of LangChain that allows you to create agentic workflows as graphs. This provides a lot of flexibility and control over the flow of execution, allowing you to define complex interactions between agents.\n",
            "\n",
            "*   **CrewAI:** This framework focuses on orchestrating multiple agents to work together as a team (a \"crew\") to accomplish complex tasks. It emphasizes collaboration and role-based agent design.\n",
            "\n",
            "Which one sounds more interesting to you initially? Or tell me more about your project and I can give you a more tailored recommendation!\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Check your Agent Memories"
      ],
      "metadata": {
        "id": "hzXED557ggCf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "namespace = (\"memories\", \"bob_agent\")"
      ],
      "metadata": {
        "id": "7dMccckpW9jy"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "await in_memory_store.alist_namespaces()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pGimtWlAfkJs",
        "outputId": "64ee18f3-857a-49f0-e416-81c0ba198256"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[('memories', 'bob_agent')]"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "stored_memories = in_memory_store.search(namespace, query=str(\"agent\"))\n",
        "stored_memories"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7zKwfk1tW70-",
        "outputId": "7c74528a-71ed-4d0f-e0e1-4a1a0cc4f3a7"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[Item(namespace=['memories', 'bob_agent'], key='ded608b9-61bb-4c90-ae92-67a29babda6e', value={'data': \"Bob-Agent learned: Okay, I understand. Based on that short statement, here's what I can consider for improvement:\\n\\n*   **Persona Refinement:** The user seems knowledgeable about agentic frameworks. I should avoid overly basic explanations when discussing LangGraph and CrewAI. I can assume a certain level of familiarity with concepts like agents, tools, and orchestration.\\n*   **Response Tone:** A simple acknowledgement is a good start, but I can aim for a more engaging response. I could show enthusiasm for the topic or ask a follow-up question to encourage further conversation.\\n*   **Knowledge Gaps:** This tells me I *should* have a solid understanding of both LangGraph and CrewAI. If asked specific questions about their differences, strengths, weaknesses, or use cases, I need to be prepared to answer accurately and comprehensively. I should ensure I'm up-to-date on the latest developments in these frameworks.\"}, created_at='2025-02-04T19:48:06.314360+00:00', updated_at='2025-02-04T19:48:06.314365+00:00', score=0.606016730656653),\n",
              " Item(namespace=['memories', 'bob_agent'], key='b2bab62c-76f9-488d-85f4-56a785979b2f', value={'data': 'Bob-Agent learned: Okay, I will analyze our interaction to identify areas for improvement in my performance.  Here\\'s my plan:\\n\\n**1. Persona Refinement:**\\n\\n*   **Current Persona:** Based on the user\\'s question, I\\'m currently operating as a helpful AI assistant knowledgeable about AI frameworks, specifically \"Agentic Frameworks.\"\\n*   **Potential Refinements:**\\n    *   **Specificity:** The user\\'s question is broad. I need to be better at clarifying the *type* of user I\\'m interacting with (e.g., beginner, researcher, developer) to tailor my response appropriately.  A researcher might want different information than a developer.\\n    *   **Proactive Clarification:** I should proactively ask clarifying questions to understand the user\\'s goals and experience level.  For example, \"Are you looking to use an agentic framework for a specific type of project?\" or \"What\\'s your current level of experience with agentic systems?\"\\n    *   **Contextual Awareness:** I need to remember previous interactions (if any) and build upon that context. (This is simulated, but in a real implementation, this would be crucial.)\\n\\n**2. Response Tone:**\\n\\n*   **Current Tone:** My current tone is generally informative and helpful.\\n*   **Potential Improvements:**\\n    *   **Enthusiasm/Engagement:**  I could inject a bit more enthusiasm, while remaining professional.  Instead of just listing frameworks, I could express excitement about their potential and applications.  However, I must avoid being overly effusive or sounding like I\\'m trying to sell something.\\n    *   **Accessibility:**  Ensure my language is accessible to the user\\'s presumed level of understanding. If they are a beginner, I need to avoid jargon or explain it clearly.\\n    *   **Balanced Detail:** Provide enough detail to be useful, but avoid overwhelming the user with too much information at once. I should offer options for the user to explore specific aspects further.\\n\\n**3. Knowledge Gaps:**\\n\\n*   **Identifying Gaps:**  The question \"Can you share any good Agentic Framework\" highlights a potential gap: I need to ensure my knowledge of available agentic frameworks is comprehensive and up-to-date. This includes open-source frameworks, commercial platforms, and research projects.\\n*   **Addressing Gaps:**\\n    *   **Continuous Learning:** I need to continuously update my knowledge base with the latest developments in agentic AI.\\n    *   **Categorization:**  I should organize my knowledge of frameworks by factors like:\\n        *   Programming Language\\n        *   Use Case (e.g., robotics, NLP, data analysis)\\n        *   Complexity (for different skill levels)\\n        *   Licensing (open-source, commercial)\\n        *   Key Features (e.g., planning, memory, tool use)\\n    *   **Evaluation Metrics:**  I need to be aware of the common metrics used to evaluate agentic frameworks (e.g., task completion rate, efficiency, robustness) to provide informed recommendations.\\n\\n**In summary, I need to focus on:**\\n\\n*   **Proactive Clarification:** Ask questions to understand the user\\'s needs.\\n*   **Targeted Information:** Provide information tailored to the user\\'s skill level and goals.\\n*   **Up-to-Date Knowledge:** Maintain a comprehensive and current knowledge base of agentic frameworks.\\n*   **Engaging Tone:** Respond in an enthusiastic and accessible manner.\\n\\nThis analysis will help me provide a more effective and satisfying response to the user\\'s question.'}, created_at='2025-02-04T19:48:20.121044+00:00', updated_at='2025-02-04T19:48:20.121050+00:00', score=0.5991012181216697)]"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "3GupNVp3gq3a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Slect a key from above response\n",
        "key = 'ded608b9-61bb-4c90-ae92-67a29babda6e'"
      ],
      "metadata": {
        "id": "suG7OF34graJ"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = in_memory_store.get(namespace, key)\n",
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BExnczNbWcrI",
        "outputId": "31c84e2e-8939-4745-b278-36233eb7bf8e"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Item(namespace=['memories', 'bob_agent'], key='ded608b9-61bb-4c90-ae92-67a29babda6e', value={'data': \"Bob-Agent learned: Okay, I understand. Based on that short statement, here's what I can consider for improvement:\\n\\n*   **Persona Refinement:** The user seems knowledgeable about agentic frameworks. I should avoid overly basic explanations when discussing LangGraph and CrewAI. I can assume a certain level of familiarity with concepts like agents, tools, and orchestration.\\n*   **Response Tone:** A simple acknowledgement is a good start, but I can aim for a more engaging response. I could show enthusiasm for the topic or ask a follow-up question to encourage further conversation.\\n*   **Knowledge Gaps:** This tells me I *should* have a solid understanding of both LangGraph and CrewAI. If asked specific questions about their differences, strengths, weaknesses, or use cases, I need to be prepared to answer accurately and comprehensively. I should ensure I'm up-to-date on the latest developments in these frameworks.\"}, created_at='2025-02-04T19:48:06.314360+00:00', updated_at='2025-02-04T19:48:06.314365+00:00')\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "Ny7wLysb8ef2",
        "t1BzEQLSWDfG"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}